package com.tuniu.agents.common.handler;

import com.tuniu.agents.advisor.ChatEnvironmentMemory;
import com.tuniu.agents.common.i18n.I18nConstants;
import com.tuniu.agents.common.message.BaseAgentNewResp;
import com.tuniu.agents.common.message.CardTypeEnum;
import com.tuniu.agents.memory.MemoryManager;
import com.tuniu.agents.message.MessageRender;
import com.tuniu.agents.tools.widgets.tool.RestaurantListTool;
import com.tuniu.agents.util.I18nUtil;
import com.tuniu.agents.util.JsonUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

@Slf4j
public class RestaurantStreamResponseHandler implements IStreamResponseHandler {

    private final String questionId;

    private final String toolName;

    private final String agentId;
    private final ChatEnvironmentMemory chatEnvironmentMemory;

    private final MemoryManager memoryManager;

    // Matching patterns
    private final Pattern bracketPattern;

    private final Pattern angleBracketPattern = Pattern.compile("<([^>]+)>");

    public static final String MESSAGE = "message";

    public static final String CONVERSATION_ID = "conversationId";

    // Pattern for matching restaurant name starting with letters followed by string (e.g., G2578)
    private final Pattern restaurantNamePattern = Pattern.compile("^([A-Za-z]\\d+)");

    private MessageRender messageRender;

    public RestaurantStreamResponseHandler(String questionId, String itemRegex, String toolName, String agentId,
                                           ChatEnvironmentMemory chatEnvironmentMemory, MemoryManager memoryManager, MessageRender messageRender) {
        if (!StringUtils.hasText(itemRegex)) itemRegex = "《([^》]+)》";

        this.bracketPattern = Pattern.compile(itemRegex);
        this.questionId = questionId;
        this.toolName = toolName;
        this.agentId = agentId;
        this.chatEnvironmentMemory = chatEnvironmentMemory;
        this.memoryManager = memoryManager;
        this.messageRender = messageRender;
    }

    public RestaurantStreamResponseHandler(String toolName, String agentId, ChatEnvironmentMemory chatEnvironmentMemory, MemoryManager memoryManager, MessageRender messageRender) {
        this("", "《([^》]+)》", toolName, agentId, chatEnvironmentMemory, memoryManager, messageRender);
    }

    @Override
    public Flux<String> handle(Map<String, Object> handlerContext, Flux<String> flux) {

        // Get tool response data
        String message = (String) handlerContext.get(MESSAGE);
        // Get conversation ID
        String conversationId = (String) handlerContext.get(CONVERSATION_ID);

        RestaurantListTool.ReqRespWrap reqRespWrap = JsonUtil.fromJson(message, RestaurantListTool.ReqRespWrap.class);

        Map<String, RestaurantListTool.Restaurant> toolResponseMap = reqRespWrap.response().data().stream().collect(Collectors.toMap(RestaurantListTool.Restaurant::name, restaurant -> restaurant, (k1, k2) -> k2));
        if (toolResponseMap.isEmpty()) {
            log.error("RestaurantStreamResponseHandler restaurantMap is empty，agentId:{},message:{}", agentId, message);
            memoryManager.setMessageRight(conversationId, Collections.singletonList(new AssistantMessage(I18nUtil.getMessage(I18nConstants.COMMON_SYSTEM_ERROR_MESSAGE_KEY), Map.of("agentId", agentId))));
            return Flux.just(BaseAgentNewResp.sendAnswerMsg(I18nUtil.getMessage(I18nConstants.COMMON_SYSTEM_ERROR_MESSAGE_KEY), questionId));
        }

        // Get card type
        String cardType = getCardTypeByToolName(toolName).getCardType();

        List<String> matchedKeys = new ArrayList<>();
        // Create a buffer accumulation stream processor
        StringBuilder buffer = new StringBuilder();
        // Store markdown formatted recommendation information
        StringBuilder markdownRestaurantRecommendations = new StringBuilder();
        AtomicInteger index = new AtomicInteger(1);

        return flux.flatMap(chunk -> {
            if (!StringUtils.hasText(chunk)) {
                return Flux.empty();
            }

            // Accumulate content
            buffer.append(chunk);
            String content = buffer.toString();

            // First try to match 《》pattern
            Matcher matcher = bracketPattern.matcher(content);
            Flux<String> response = processMatches(matcher, buffer, toolResponseMap, cardType, matchedKeys, message, conversationId, markdownRestaurantRecommendations, index);
            if (response != null) {
                return response;
            }

            // If no 《》match found, try matching <> pattern
            matcher = angleBracketPattern.matcher(content);
            response = processMatches(matcher, buffer, toolResponseMap, cardType, matchedKeys, message, conversationId, markdownRestaurantRecommendations, index);
            if (response != null) {
                return response;
            }

            return Flux.empty();
        }).concatWith(Mono.fromCallable(() -> {
            String remainingContent = buffer.toString();

            // Add suggestions and summary content
            if (StringUtils.hasText(remainingContent)) {
                markdownRestaurantRecommendations.append(remainingContent.trim());
            }
            // Store markdown formatted recommendations in context
            if (!markdownRestaurantRecommendations.isEmpty()) {
                memoryManager.setMessageRight(conversationId, Collections.singletonList(new AssistantMessage(markdownRestaurantRecommendations.toString(), Map.of("agentId", agentId))));
            }
            return StringUtils.hasText(remainingContent) ?
                    BaseAgentNewResp.sendAnswerMsg(null, remainingContent) : "";
        })).filter(StringUtils::hasText);
    }

    /**
     * Extract key from matcher
     */
    private String extractKey(Matcher matcher) {
        return matcher.group(1).trim();
    }

    /**
     * Extract recommendation reason
     */
    private String extractReason(String content, int endIndex) {
        return content.substring(0, endIndex).trim();
    }

    private CardTypeEnum getCardTypeByToolName(String toolName) {
        if (toolName.equals(RestaurantListTool.TOOLNAME)) {
            return CardTypeEnum.RESTAURANT;
        }
        return CardTypeEnum.RESTAURANT;
    }

    /**
     * Process matching results
     * @return Returns response string if match is successful, otherwise returns null
     */
    private Flux<String> processMatches(Matcher matcher, StringBuilder buffer, Map<String, RestaurantListTool.Restaurant> toolResponseMap,
                                        String cardType, List<String> matchedKeys, String message, String conversationId, StringBuilder markdownRestaurantRecommendations, AtomicInteger index) {
        if (matcher.find()) {
            // Extract restaurant name
            String key = extractKey(matcher);

            int endIndex = matcher.start();
            // Recommendation reason
            String reason = extractReason(buffer.toString(), endIndex);

            matchedKeys.add(key);

            // Update buffer, remove processed part
            buffer.delete(0, matcher.end());

            // Get restaurant by restaurant name
            RestaurantListTool.Restaurant restaurant = toolResponseMap.get(key);

            if (restaurant == null) {
                Map.Entry<String, RestaurantListTool.Restaurant> stringRestaurantEntry = toolResponseMap.entrySet().stream().filter(e -> e.getKey().contains(key)).findFirst().orElse(null);
                if (stringRestaurantEntry != null) {
                    restaurant = stringRestaurantEntry.getValue();
                }
                if (restaurant == null) {
                    log.error("RestaurantStreamResponseHandler toolResponseMap.get(key) is empty，key:{},toolResponseMap:{}", key, JsonUtil.toJson(toolResponseMap));
                    return null;
                }
            }

            // Record markdown information to be stored
            String saveInMdReason = !reason.contains("：") ? reason : reason.substring(reason.indexOf("：") + 1);
            saveInMdReason = !saveInMdReason.contains(":") ? saveInMdReason : saveInMdReason.substring(saveInMdReason.indexOf(":") + 1);

            Map<String, Object> restaurantDataMap = JsonUtil.toMap(JsonUtil.toJson(restaurant));
            restaurantDataMap.put("recommendationReason", saveInMdReason);
            // Markdown message
            String markdownResponse = null;
            if (CardTypeEnum.RESTAURANT.getCardType().equals(cardType)) {
                String acceptLanguage = I18nUtil.getLocaleLanguageTag(conversationId);
                markdownResponse = messageRender.renderTemplate("widgets/restaurant-list-one.ftlh", Map.of(I18nConstants.LANGUAGE_TAG, acceptLanguage), restaurantDataMap);
            }

            List<Object> data = new ArrayList<>();
            data.add(new Object());
            String response = BaseAgentNewResp.sendCardMsg(cardType, data, markdownResponse, null);
            log.info("RestaurantStreamResponseHandler handle message:{},response:{}", message, response);
            return Flux.just(response);

        }
        return null;
    }

}